#include <winsock2.h>
#include <ws2tcpip.h>
#include <iostream>
#include <thread>
#include <vector>
#include <stdio.h>
#include <string>
#include <regex>
#pragma comment(lib, "ws2_32.lib")
using namespace std;

#define BUFFER_SIZE 4096

// 处理客户端连接的函数
void handle_client(SOCKET clientSock);
// 将域名转换为ip的函数
string resolve_domain_to_ip(const  string& domain);

int main(int argc, char* argv[]) {
    if (argc != 2) {
         cerr << "Usage: " << argv[0] << " <port>" << endl;
        return 1;
    }

    // 初始化 WinSock，启动网络通信功能
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        cerr << "Failed to initialize WinSock!" << endl;
        return 1;
    }

    // 创建一个监听套接字，用于接收客户端连接
    SOCKET listenSock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (listenSock == INVALID_SOCKET) {
        cerr << "Failed to create socket!" << endl;
        WSACleanup();
        return 1;
    }

    // 配置服务器地址结构
    sockaddr_in serverAddr;
    serverAddr.sin_family = AF_INET; // 地址族：IPv4
    serverAddr.sin_port = htons(atoi(argv[1])); // 将命令行输入的端口号转换为网络字节序
    serverAddr.sin_addr.s_addr = INADDR_ANY;    // 监听所有可用的网络接口地址

    // 将套接字绑定到指定的 IP 地址和端口
    if (bind(listenSock, (sockaddr*)&serverAddr, sizeof(serverAddr)) == SOCKET_ERROR) {
        cerr << "Failed to bind socket!" << endl;
        closesocket(listenSock);
        WSACleanup();
        return 1;
    }

    // 开始监听，准备接受客户端连接
    if (listen(listenSock, SOMAXCONN) == SOCKET_ERROR) {
        cerr << "Failed to listen on socket!" << endl;
        closesocket(listenSock);
        WSACleanup();
        return 1;
    }

    cout << "SOCKS proxy server is running on port " << argv[1] << endl;

    // 等待客户端连接并处理
    while (true) {
        sockaddr_in clientAddr; // 用于存储客户端地址信息
        int clientAddrLen = sizeof(clientAddr);

        // 接受客户端连接
        SOCKET clientSock = accept(listenSock, (sockaddr*)&clientAddr, &clientAddrLen);
        if (clientSock == INVALID_SOCKET) {
            cerr << "Failed to accept client connection!" << endl;
            continue;
        }

        // 打印客户端的 IP 和端口信息
        cout << "Accepted connection from "
                  << inet_ntoa(clientAddr.sin_addr) << ":" << ntohs(clientAddr.sin_port) << endl;

        // 为每个客户端连接创建一个独立的线程进行处理
        thread(handle_client, clientSock).detach();
    }
}

// 处理客户端连接的具体逻辑
void handle_client(SOCKET clientSock) {
    char buffer[BUFFER_SIZE]; // 数据缓冲区
    int bytesRead; // 接收数据的字节数

    // 接收 SOCKS 版本号
    bytesRead = recv(clientSock, buffer, BUFFER_SIZE, 0);
    if (bytesRead <= 0) {
        cerr << "Failed to read from client!" << endl;
        closesocket(clientSock);
        return;
    }
    // 根据 SOCKS 版本号分配处理逻辑
    if (buffer[0] == 4) {
        cout << "SOCKS4 request received." << endl;
    }
    else if (buffer[0] == 5) {
        cout << "SOCKS5 request received." << endl;

        // SOCKS5 的认证阶段：通知客户端无需认证
        send(clientSock, "\x05\x00", 2, 0);

        // 接收客户端的连接请求
        bytesRead = recv(clientSock, buffer, BUFFER_SIZE, 0);
        if (bytesRead <= 0) {
            cerr << "Failed to read connection request!" << endl;
            closesocket(clientSock);
            return;
        }

        if(buffer[bytesRead - 1] != -69) {
            for(int i = 0; i < bytesRead; i++) {
                printf("%d ", buffer[i]);
            }
            cout << endl;
        }

        if (buffer[1] == 0x01) {
            sockaddr_in targetAddr; // 目标服务器地址
            targetAddr.sin_family = AF_INET;

            // 判断目标地址是 IP 还是域名
            string targetHost;
            if (buffer[3] == 1) {// 如果目标地址是 IP 地址
                targetAddr.sin_port = *(u_short*)&buffer[4];
                targetAddr.sin_addr.s_addr = *(u_long*)&buffer[8];
                targetHost = inet_ntoa(targetAddr.sin_addr);
            }
            if (buffer[3] == 0x03) { // 如果地址类型是域名
                uint8_t domainLength = buffer[4];
                targetHost = string(reinterpret_cast<char*>(&buffer[5]), domainLength);
                uint16_t port = *(uint16_t*)&buffer[5 + domainLength];
                cout << targetHost << endl;
                // 获取解析的 IP 地址
                string ip = resolve_domain_to_ip(targetHost);
                if (ip.empty()) {
                    cerr << "Failed to resolve domain: " << targetHost << endl;
                    closesocket(clientSock);
                    return;
                }
                targetAddr.sin_addr.s_addr = inet_addr(ip.c_str()); // 转换为 IP 地址
                targetAddr.sin_port = port; // 设置目标端口
            }

            // 创建一个新套接字连接目标服务器
            SOCKET targetSock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
            if (connect(targetSock, (sockaddr*)&targetAddr, sizeof(targetAddr)) == SOCKET_ERROR) {
                cerr << "Failed to connect to target server!" << endl;
                closesocket(clientSock);
                return;
            }

            // 返回成功响应给客户端
            send(clientSock, "\x05\x00\x00\x01\x00\x00\x00\x00\x00\x00", 10, 0);

            // 转发数据
            fd_set fds;
            while (true) {
                FD_ZERO(&fds);
                FD_SET(clientSock, &fds);
                FD_SET(targetSock, &fds);

                int maxFd = max(clientSock, targetSock) + 1;

                // 使用 select 监控套接字上的活动
                int activity = select(maxFd, &fds, nullptr, nullptr, nullptr);
                if (activity <= 0) break;

                // 如果客户端发送了数据，转发给目标服务器
                if (FD_ISSET(clientSock, &fds)) {
                    bytesRead = recv(clientSock, buffer, BUFFER_SIZE, 0);
                    if (bytesRead <= 0) break; // 客户端关闭连接
                    send(targetSock, buffer, bytesRead, 0);
                }

                // 如果目标服务器发送了数据，转发给客户端
                if (FD_ISSET(targetSock, &fds)) {
                    bytesRead = recv(targetSock, buffer, BUFFER_SIZE, 0);
                    if (bytesRead <= 0) break;
                    send(clientSock, buffer, bytesRead, 0);
                }
            }
            closesocket(targetSock);
        }
    }
    closesocket(clientSock);
}

// 将域名解析为 IP 地址
string resolve_domain_to_ip(const string& domain) {
    addrinfo hints, *res;
    int err;
    char ip[INET6_ADDRSTRLEN];
    memset(&hints, 0, sizeof(hints));
    hints.ai_family = AF_INET;

    err = getaddrinfo(domain.c_str(), nullptr, &hints, &res);
    if (err != 0) {
        return "";
    }

    // 提取 IP 地址
    sockaddr_in* ipv4 = (struct sockaddr_in*)res->ai_addr;
    inet_ntop(AF_INET, &ipv4->sin_addr, ip, sizeof(ip));

    freeaddrinfo(res);
    return string(ip);
}
